import os
import json
import pandas as pd
import numpy as np
from flask import Flask, request, jsonify
from flask_cors import CORS
from flask_jwt_extended import JWTManager, create_access_token, jwt_required, get_jwt_identity
from datetime import timedelta
import matplotlib.pyplot as plt
import io
import base64
from werkzeug.utils import secure_filename

# 导入分析模块
try:
    from icaMode import ICAAnalyzer, train_ica_model, get_reconstruction_error
except ImportError:
    print("无法导入ICA模块，请确保icaMode.py文件存在")

try:
    from aeMode import AEAnalyzer, train_ae_model
except ImportError:
    print("无法导入AE模块，请确保aeMode.py文件存在")

# 初始化Flask应用
app = Flask(__name__)
CORS(app)  # 允许跨域请求

# JWT配置
app.config['JWT_SECRET_KEY'] = 'huatai-intel-diagnosis-system'  # 生产环境中应使用环境变量
app.config['JWT_ACCESS_TOKEN_EXPIRES'] = timedelta(hours=12)
jwt = JWTManager(app)

# 创建上传文件夹
UPLOAD_FOLDER = 'uploads'
if not os.path.exists(UPLOAD_FOLDER):
    os.makedirs(UPLOAD_FOLDER)
app.config['UPLOAD_FOLDER'] = UPLOAD_FOLDER

# 用户数据（实际项目中应使用数据库）
users = {
    "admin": {
        "password": "admin",
        "role": "admin"
    }
}

# 存储分析结果的字典
analysis_results = {}

@app.route('/api/login', methods=['POST'])
def login():
    """用户登录API"""
    data = request.get_json()
    username = data.get('username', '')
    password = data.get('password', '')
    
    if username in users and users[username]['password'] == password:
        # 创建JWT令牌
        access_token = create_access_token(identity=username)
        return jsonify(access_token=access_token, username=username), 200
    else:
        return jsonify(message="用户名或密码错误"), 401

@app.route('/api/upload', methods=['POST'])
@jwt_required()
def upload_file():
    """上传数据文件API"""
    current_user = get_jwt_identity()
    
    # 检查是否有文件
    if 'file' not in request.files:
        return jsonify(message="没有文件"), 400
    
    file = request.files['file']
    
    # 检查文件名
    if file.filename == '':
        return jsonify(message="没有选择文件"), 400
    
    # 检查文件类型
    if file and (file.filename.endswith('.xlsx') or file.filename.endswith('.csv')):
        filename = secure_filename(file.filename)
        file_path = os.path.join(app.config['UPLOAD_FOLDER'], filename)
        file.save(file_path)
        
        # 读取数据
        try:
            if filename.endswith('.xlsx'):
                data_df = pd.read_excel(file_path)
            else:
                data_df = pd.read_csv(file_path, header=None)
                data_df.columns = [f'X{i+1}' for i in range(data_df.shape[1])]
            
            # 存储数据
            user_key = f"user_{current_user}"
            if user_key not in analysis_results:
                analysis_results[user_key] = {}
            
            analysis_results[user_key]['data'] = data_df
            analysis_results[user_key]['file_path'] = file_path
            
            # 返回数据的基本信息
            return jsonify({
                "message": "文件上传成功",
                "rows": data_df.shape[0],
                "columns": data_df.shape[1],
                "column_names": data_df.columns.tolist()
            }), 200
        except Exception as e:
            return jsonify(message=f"处理文件时出错: {str(e)}"), 500
    else:
        return jsonify(message="不支持的文件类型，仅支持.xlsx和.csv文件"), 400

@app.route('/api/train/ica', methods=['POST'])
@jwt_required()
def train_ica():
    """训练ICA模型API"""
    current_user = get_jwt_identity()
    user_key = f"user_{current_user}"
    
    # 检查是否有上传的数据
    if user_key not in analysis_results or 'data' not in analysis_results[user_key]:
        return jsonify(message="请先上传数据"), 400
    
    # 获取请求数据
    data = request.get_json()
    n_components = data.get('n_components', None)
    
    if not n_components:
        return jsonify(message="缺少组件数量参数"), 400
    
    try:
        n_components = int(n_components)
        data_df = analysis_results[user_key]['data']
        
        # 检查组件数量是否有效
        if n_components < 2 or n_components > data_df.shape[1]:
            return jsonify(message=f"组件数量必须在2到{data_df.shape[1]}之间"), 400
        
        # 创建ICA分析器
        analyzer = ICAAnalyzer()
        processed_data = analyzer.preprocess_data(data_df)
        
        if processed_data is None:
            return jsonify(message="数据预处理失败"), 400
        
        # 训练ICA模型
        analyzer, figures = train_ica_model(processed_data, n_components=n_components)
        
        if analyzer is None:
            return jsonify(message="模型训练失败"), 400
        
        # 存储分析结果
        analysis_results[user_key]['ica_analyzer'] = analyzer
        
        # 转换图表为Base64
        figure_data = []
        for i, fig in enumerate(figures):
            buf = io.BytesIO()
            fig.savefig(buf, format='png', dpi=100, bbox_inches='tight')
            buf.seek(0)
            img_str = base64.b64encode(buf.read()).decode('utf-8')
            figure_data.append({
                'id': i,
                'title': f'ICA分析图 {i+1}',
                'image_data': img_str
            })
            plt.close(fig)  # 关闭图形以释放内存
        
        # 返回成功消息和图表数据
        return jsonify({
            "message": "ICA模型训练成功",
            "figures": figure_data,
            "n_components": n_components
        }), 200
    except Exception as e:
        return jsonify(message=f"训练模型时出错: {str(e)}"), 500

@app.route('/api/analysis/ica/contribution', methods=['GET'])
@jwt_required()
def get_ica_contribution():
    """获取ICA贡献图API"""
    current_user = get_jwt_identity()
    user_key = f"user_{current_user}"
    
    # 检查是否有训练好的模型
    if (user_key not in analysis_results or 
        'ica_analyzer' not in analysis_results[user_key]):
        return jsonify(message="请先训练ICA模型"), 400
    
    try:
        analyzer = analysis_results[user_key]['ica_analyzer']
        
        # 获取异常点
        anomaly_indices = np.where(analyzer.anomalies)[0]
        if len(anomaly_indices) == 0:
            return jsonify(message="未检测到异常点"), 200
        
        # 生成贡献图
        contribution_figures = []
        for i, idx in enumerate(anomaly_indices):
            figs = analyzer.plot_contribution_for_anomalies([idx])
            for j, fig in enumerate(figs):
                buf = io.BytesIO()
                fig.savefig(buf, format='png', dpi=100, bbox_inches='tight')
                buf.seek(0)
                img_str = base64.b64encode(buf.read()).decode('utf-8')
                contribution_figures.append({
                    'id': i * len(figs) + j,
                    'anomaly_index': int(idx),
                    'title': f'异常点 {idx} 贡献图',
                    'image_data': img_str
                })
                plt.close(fig)  # 关闭图形以释放内存
        
        # 返回贡献图数据
        return jsonify({
            "message": "贡献图生成成功",
            "anomaly_count": len(anomaly_indices),
            "figures": contribution_figures
        }), 200
    except Exception as e:
        return jsonify(message=f"生成贡献图时出错: {str(e)}"), 500

@app.route('/api/analysis/ica/diagnosis', methods=['GET'])
@jwt_required()
def get_ica_diagnosis():
    """获取ICA诊断报告API"""
    current_user = get_jwt_identity()
    user_key = f"user_{current_user}"
    
    # 检查是否有训练好的模型
    if (user_key not in analysis_results or 
        'ica_analyzer' not in analysis_results[user_key]):
        return jsonify(message="请先训练ICA模型"), 400
    
    try:
        analyzer = analysis_results[user_key]['ica_analyzer']
        
        # 生成诊断报告
        report = analyzer.diagnose_faults()
        
        # 返回诊断报告
        return jsonify({
            "message": "诊断报告生成成功",
            "report": report
        }), 200
    except Exception as e:
        return jsonify(message=f"生成诊断报告时出错: {str(e)}"), 500

@app.route('/api/train/ae', methods=['POST'])
@jwt_required()
def train_ae():
    """训练AE模型API"""
    current_user = get_jwt_identity()
    user_key = f"user_{current_user}"
    
    # 检查是否有上传的数据
    if user_key not in analysis_results or 'data' not in analysis_results[user_key]:
        return jsonify(message="请先上传数据"), 400
    
    # 获取请求数据
    data = request.get_json()
    encoding_dim = data.get('encoding_dim', None)
    epochs = data.get('epochs', 100)
    
    if not encoding_dim:
        return jsonify(message="缺少编码维度参数"), 400
    
    try:
        encoding_dim = int(encoding_dim)
        epochs = int(epochs)
        data_df = analysis_results[user_key]['data']
        
        # 创建AE分析器
        analyzer = AEAnalyzer()
        
        # 训练AE模型
        analyzer, model_summary, loss_history, reconstruction_error = train_ae_model(
            data_df, encoding_dim=encoding_dim, epochs=epochs
        )
        
        if analyzer is None:
            return jsonify(message="模型训练失败"), 400
        
        # 存储分析结果
        analysis_results[user_key]['ae_analyzer'] = analyzer
        
        # 创建损失历史图
        plt.figure(figsize=(10, 6))
        plt.plot(loss_history)
        plt.title('Model Loss During Training')
        plt.ylabel('Loss')
        plt.xlabel('Epoch')
        plt.grid(True)
        
        buf = io.BytesIO()
        plt.savefig(buf, format='png', dpi=100, bbox_inches='tight')
        buf.seek(0)
        loss_img_str = base64.b64encode(buf.read()).decode('utf-8')
        plt.close()
        
        # 返回成功消息和图表数据
        return jsonify({
            "message": "AE模型训练成功",
            "model_summary": model_summary,
            "loss_figure": loss_img_str,
            "encoding_dim": encoding_dim,
            "epochs": epochs,
            "final_loss": float(loss_history[-1]) if loss_history else None,
            "mean_reconstruction_error": float(np.mean(reconstruction_error)) if reconstruction_error is not None else None
        }), 200
    except Exception as e:
        return jsonify(message=f"训练模型时出错: {str(e)}"), 500

if __name__ == '__main__':
    app.run(debug=True, port=5000) 